#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python常用模块库 
@File ：websocket_protocol.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/2/20 上午12:16 
"""

import asyncio
import logging
from typing import Callable, Any

import websockets


class WebSocketProtocol:
	def __init__(self, host: str, port: int, callback: Callable[[bytes], Any] = None):
		self.host = host
		self.port = port
		self.callback = callback
		self.server = None

	async def _handler(self, websocket, path):
		"""WebSocket 处理器"""
		async for message in websocket:
			if self.callback:
				self.callback(message.encode())

	def start(self):
		"""启动WebSocket服务器"""

		async def _start():
			self.server = await websockets.serve(self._handler, self.host, self.port)
			logging.info(f"WebSocket server started on ws://{self.host}:{self.port}")
			await self.server.wait_closed()

		asyncio.get_event_loop().run_until_complete(_start())

	def stop(self):
		"""停止WebSocket服务器"""
		if self.server:
			self.server.close()
			logging.info("WebSocket server stopped.")

	async def send(self, data: bytes) -> Any:
		"""发送WebSocket数据"""
		async with websockets.connect(f"ws://{self.host}:{self.port}") as websocket:
			await websocket.send(data)
			return await websocket.recv()
